functor OptOnAst (S : OPT_ON_AST_STRUCTS) : OPT_ON_AST =
struct
  open S
  
  structure A = Ast
  structure ErrorMsg = ErrorMsg
  
  local open A
  in structure Id = A.Id
     structure Intl = A.Intl
     structure B = A.Binop
     structure T = A.Tipe
     structure L = A.Le
     structure S = A.Structs
     structure M = A.Stm
     structure F = A.Function
     structure P = A.Program
  end

  val error = ErrorMsg.error

  val tag = false
  fun debugPrint s = 
      if tag
      then print (String.concat [s, "\n"])
      else ()

structure Cf =
struct
  fun shift b = if b then L.True else L.False
  
  fun cfOneExp exp =
      case exp
       of L.Expvoid => L.Expvoid
        | L.IntNum i => L.IntNum i
        | L.True => L.True
        | L.False => L.False
        | L.Null => L.Null
        | L.Lval lval => L.Lval (cfLval lval)
        (*--------------------------- + ------------------------*)
        | L.BinopExp (L.IntNum i, B.Plus, L.IntNum j, pos) =>
          L.IntNum (Intl.plus (i, j))
        | L.BinopExp (e1, B.Plus, e2, pos) =>
          L.BinopExp (cfOneExp e1, B.Plus, cfOneExp e2, pos)
        (*--------------------------- - ------------------------*)
        | L.BinopExp(L.IntNum i, B.Minus, L.IntNum j, pos) =>
          L.IntNum (Intl.minus (i, j))
        | L.BinopExp (e1, B.Minus, e2, pos) =>
          L.BinopExp (cfOneExp e1, B.Minus, cfOneExp e2, pos)
        (*--------------------------- * ------------------------*)
        | x as L.BinopExp (L.IntNum i, B.Times, e2, pos) =>
          if Intl.isZero i
          then L.IntNum (Intl.zero)
          else (case e2
                of L.IntNum j => L.IntNum (Intl.times (i, j))
                 | _ => x)
        | x as L.BinopExp (e1, B.Times, L.IntNum j, pos) =>
          if Intl.isZero j
          then L.IntNum (Intl.zero)
          else x
        | L.BinopExp (e1, B.Times, e2, pos) =>
          L.BinopExp (cfOneExp e1, B.Times, cfOneExp e2, pos)
        (*--------------------------- / ------------------------*)
        | x as L.BinopExp (L.IntNum i, B.Divide, e2, pos) =>
          if Intl.isZero i
          then L.IntNum (Intl.zero)
          else (case e2
                of L.IntNum j =>
                   if Intl.isZero j
                   then (error pos "divide by zero"
                       ; L.IntNum (Intl.zero))
                   else L.IntNum (Intl.divide (i, j))
                 | _ => x)
        | L.BinopExp (e1, B.Divide, e2, pos) =>
          L.BinopExp(cfOneExp e1, B.Divide, cfOneExp e2, pos)
        (*--------------------------- % ------------------------*)
        | x as L.BinopExp(L.IntNum i, B.Modus, e2, pos) =>
          if Intl.isZero i
          then L.IntNum (Intl.zero)
          else (case e2
                of L.IntNum j =>
                   if Intl.isZero j
                   then (error pos "divided by zero"
                       ; L.IntNum (Intl.zero))
                   else L.IntNum (Intl.modus (i, j))
                 | _ => x)
        | L.BinopExp (e1, B.Modus, e2, pos) => 
          L.BinopExp (cfOneExp e1, B.Modus, cfOneExp e2, pos)
        (*-------------------------- && ------------------------*)
        | L.BinopExp(L.True, B.And, e2, pos) => e2
        | L.BinopExp(L.False, B.And, e2, pos) => L.False
        (* WRONG!!!
        | L.BinopExp(e1, B.And, L.True, pos)) = e2
        | L.BinopExp(e1, B.And, L.False, pos)) = L.False*)
        | L.BinopExp(e1, B.And, e2, pos) =>
          L.BinopExp(cfOneExp e1, B.And, cfOneExp e2, pos)
        (*--------------------------- || ------------------------*)
        | L.BinopExp (L.True, B.Or, e2, pos) => L.True
        | L.BinopExp (L.False, B.Or, e2, pos) => e2
        | L.BinopExp(e1, B.Or, e2, pos) => 
          L.BinopExp(cfOneExp e1, B.Or, cfOneExp e2, pos)
        (*--------------------------- == ------------------------*)
        | L.BinopExp(L.IntNum(i), B.Eq, L.IntNum(j), pos) => shift (Intl.equals (i, j))
        | L.BinopExp(e1, B.Eq, e2, pos) =>
          L.BinopExp(cfOneExp e1, B.Eq, cfOneExp e2, pos)
        (*--------------------------- != ------------------------*)
        | L.BinopExp(L.IntNum i, B.Ne, L.IntNum j, pos) => shift (Intl.ne (i, j))
        | L.BinopExp(e1, B.Ne, e2, pos) =>
          L.BinopExp(cfOneExp e1, B.Ne, cfOneExp e2, pos)
        (*--------------------------- > -------------------------*)
        | L.BinopExp(L.IntNum(i), B.Gt, L.IntNum(j), pos) => shift (Intl.gt (i, j))
        | L.BinopExp(e1, B.Gt, e2, pos) =>
          L.BinopExp(cfOneExp e1, B.Gt, cfOneExp e2, pos) before
          debugPrint ">>>>>>>>>>>>>>>"
        (*--------------------------- >= ------------------------*)
        | L.BinopExp(L.IntNum(i), B.Ge, L.IntNum(j), pos) => shift (Intl.ge (i, j))
        | L.BinopExp(e1, B.Ge, e2, pos) =>
          L.BinopExp(cfOneExp e1, B.Ge, cfOneExp e2, pos)
        (*-------------------------- < -------------------------*)
        | L.BinopExp(L.IntNum(i), B.Lt, L.IntNum(j), pos) => shift (Intl.lt (i, j))
        | L.BinopExp(e1, B.Lt, e2, pos) =>
          L.BinopExp(cfOneExp e1, B.Lt, cfOneExp e2, pos)
        (*-------------------------- <= ------------------------*)
        | L.BinopExp (L.IntNum(i), B.Le, L.IntNum(j), pos) => shift (Intl.le (i, j))
        | L.BinopExp(e1, B.Le, e2, pos) =>
          L.BinopExp(cfOneExp e1, B.Le, cfOneExp e2, pos)
        (*-------------------------- not ---------------------*)
        | L.NotExp (L.True, pos) => L.False
        | L.NotExp (L.False, pos) => L.True
        | L.NotExp (e, pos) => L.NotExp (cfOneExp e, pos)
        | L.Result => L.Result
        (*--------------------------- imply --------------------*)
        | L.Imply (L.True, e2, pos) => e2
        | L.Imply (L.False, e2, pos) => L.True
        | L.Imply (e1, L.True, pos) => L.True
        | L.Imply (e1, L.False, pos) => L.NotExp (e1, pos)
        | L.Imply (e1, e2, pos) => L.Imply (cfOneExp e1, cfOneExp e2, pos)
        | L.Iff (e1, e2, pos) => L.BinopExp (L.Imply (e1, e2, pos), B.And, L.Imply (e2, e1, pos), pos)
        | L.Forall (id, ty, e, pos) => L.Forall (id, ty, e, pos)
        | L.Exists (id, ty, e, pos) => L.Exists (id, ty, e, pos)
        | L.Pi (lss, pos) => L.Pi (lss, pos)
        | L.N (ls, pos) => L.N (ls, pos)
        | L.D (ls, pos) => L.D (ls, pos)
        | L.Prev (lval, pos) => L.Prev (lval, pos)
        | L.LList (exp, pos) => L.LList (exp, pos)
        | L.Tree (exp, pos) => L.Tree (exp, pos)
        | L.Pred (id, exp, pos) => L.Pred (id, exp, pos)
        | L.Let (e1, e2, e3, pos) => L.Let (e1, e2, e3, pos)
      
  and cfExp exp =
      let fun loop e =
              let val res = cfOneExp e
              in  if L.equalsExp (res, e)
                  then res
                  else (debugPrint "loop in exp optimization"
                      ; loop res)
              end
      in  loop exp
      end
      
  and cfLval lval = lval
  
  and mkSomeList x = SOME [x]
  
  and mkSome x = SOME x
  
  and cfStm (stm : M.t) : M.t list option =
      case stm
       of M.AsnExp (lval, exp, pos) =>
          let val new = M.AsnExp (cfLval lval, cfExp exp, pos)
          in  mkSomeList new
          end 
        | M.AsnCall (lval, id, es, pos) =>
          let val new = M.AsnCall (lval, id, List.map cfExp es, pos)
          in  mkSomeList new
          end
        | M.AsnAlloc (lval, ty, pos) =>
          let val new = M.AsnAlloc (cfLval lval, ty, pos)
          in  mkSomeList new
          end
        (*---------------------------- if ---------------------------------------------*)
        | M.If (L.True, s1, s2, pos) => mkSome s1
        | M.If (L.False, s1, s2, pos) => mkSome s2
        | M.If (e, s1, s2, pos) =>
          let val new = M.If (cfExp e, cfStms s1, cfStms s2, pos)
          in  mkSomeList new
          end
        (*---------------------------- while ------------------------------------------*)
        | M.While {condition = L.False, invariant, body, pos} => NONE
        | M.While {condition = L.True, invariant, body, pos} =>
          let val new = M.While {invariant = invariant,
                                 condition = L.True,
                                 body = cfStms body,
                                 pos = pos}
          in  mkSomeList new
          end
        | M.While {invariant, condition, body, pos} =>
          let val new = M.While {invariant = invariant,
                                 condition = cfExp condition,
                                 body = cfStms body,
                                 pos = pos}
          in  debugPrint "loop in while"
            ; mkSomeList new
          end
        | M.Return (e, pos) =>
          let val new = M.Return (cfExp e, pos)
          in  mkSomeList new
          end
        | M.Call (id, es, pos) =>
          let val new = M.Call (id, List.map cfExp es, pos) 
          in  mkSomeList new
          end
        | M.Free (e, pos) =>
          let val new = M.Free (cfExp e, pos)
          in  mkSomeList new
          end
        | M.Prop (e, pos) =>
          let val new = M.Prop (cfExp e, pos)
          in  mkSomeList new
          end
  
  and clean (l : M.t list option list) : M.t list =
      let fun f l =
              case l
               of [] => []
                | (NONE :: xs) => f xs
                | (SOME x :: xs) => x :: (f xs)
      in  List.concat (f l)
      end
      
  and cfStms (ss : M.t list) : M.t list = clean (List.map cfStm ss)
     
  fun cfTopStms stms = 
      let fun loop s =
              let val res = cfStms s
              in  if M.equalsStms (s, res)
                  then res
                  else (debugPrint "loop in stm optimization again"
                       ; loop res)
              end
      in  loop stms
      end
      
  fun cfFunction f =
      case f
       of F.T {name,
               rettype,
               args,
               locals,
               stm,
               pre,
               post,
               pos} => 
          F.T {name = name,
               rettype = rettype,
               args = args,
               locals = locals,
               stm = cfTopStms stm,
               pre = pre,
               post = post,
               pos = pos}
      
  fun cfProgram (p : P.t) : P.t =
      case p
       of P.T {structs,
               gvars,
               funcs,
               pos} =>
          P.T {structs = structs,
               gvars = gvars,
               funcs = List.map cfFunction funcs,
               pos = pos}
  
  val top = cfProgram
end

structure DeadDec =
struct
  structure Status :
  sig
    datatype t
      = Assigned
      | Used      
  end =
  struct
    datatype t
      = Assigned
      | Used
  end
  
  structure Env = BinaryMapFn (type ord_key = Id.t
                               val compare = Id.compare)
  
  fun empty () : Status.t Env.map ref = ref Env.empty
  
  val genv = empty ()
  
  fun enter (x, k) =
      let val new = Env.insert (!genv, x, k)
      in  genv := new
      end
      
  fun look (x) =
      let val k = Env.find (!genv, x)
      in  k
      end
  
  (****************************************************************************************)
  fun cfOneExp exp =
      case exp
       of L.Expvoid => L.Expvoid
        | L.IntNum i => L.IntNum i
        | L.True => L.True
        | L.False => L.False
        | L.Null => L.Null
        | L.Lval lval => L.Lval (cfLval lval)
        | L.BinopExp (e1, b, e2, pos) => raise Fail "to do"
        | L.NotExp (e, pos) => raise Fail "to do"
        | L.Result => L.Result
        | L.Imply (e, e2, pos) => raise Fail "to do"
        | L.Iff (e1, e2, pos) => L.BinopExp (L.Imply (e1, e2, pos), B.And, L.Imply (e2, e1, pos), pos)
        | L.Forall (id, ty, e, pos) => L.Forall (id, ty, e, pos)
        | L.Exists (id, ty, e, pos) => L.Exists (id, ty, e, pos)
        | L.Pi (lss, pos) => L.Pi (lss, pos)
        | L.N (ls, pos) => L.N (ls, pos)
        | L.D (ls, pos) => L.D (ls, pos)
        | L.Prev (lval, pos) => L.Prev (lval, pos)
        | L.LList (exp, pos) => L.LList (exp, pos)
        | L.Tree (exp, pos) => L.Tree (exp, pos)
        | L.Pred (id, exp, pos) => L.Pred (id, exp, pos)
        | L.Let (e1, e2, e3, pos) => L.Let (e1, e2, e3, pos)
      
  and cfExp exp =
      let fun loop e =
              let val res = cfOneExp e
              in  if L.equalsExp (res, e)
                  then res
                  else (debugPrint "loop in exp optimization"
                      ; loop res)
              end
      in  loop exp
      end
      
  and cfLval lval = lval
  
  and mkSomeList x = SOME [x]
  
  and mkSome x = SOME x
  
  and cfStm (stm : M.t) : M.t list option =
      case stm
       of M.AsnExp (lval, exp, pos) =>
          let val new = M.AsnExp (cfLval lval, cfExp exp, pos)
          in  mkSomeList new
          end 
        | M.AsnCall (lval, id, es, pos) =>
          let val new = M.AsnCall (lval, id, List.map cfExp es, pos)
          in  mkSomeList new
          end
        | M.AsnAlloc (lval, ty, pos) =>
          let val new = M.AsnAlloc (cfLval lval, ty, pos)
          in  mkSomeList new
          end
        (*---------------------------- if ---------------------------------------------*)
        | M.If (L.True, s1, s2, pos) => mkSome s1
        | M.If (L.False, s1, s2, pos) => mkSome s2
        | M.If (e, s1, s2, pos) =>
          let val new = M.If (cfExp e, cfStms s1, cfStms s2, pos)
          in  mkSomeList new
          end
        (*---------------------------- while ------------------------------------------*)
        | M.While {invariant, condition = L.False, body, pos} => NONE
        | M.While {invariant, condition = L.True, body, pos} =>
          let val new = M.While {invariant = invariant,
                                 condition = L.True,
                                 body = cfStms body,
                                 pos = pos}
          in  mkSomeList new
          end
        | M.While {invariant, condition, body, pos} =>
          let val new = M.While {invariant = invariant, 
                                 condition = cfExp condition,
                                 body = cfStms body,
                                 pos = pos}
          in  debugPrint "loop in while"
            ; mkSomeList new
          end
        | M.Return (e, pos) =>
          let val new = M.Return (cfExp e, pos)
          in  mkSomeList new
          end
        | M.Call (id, es, pos) =>
          let val new = M.Call (id, List.map cfExp es, pos) 
          in  mkSomeList new
          end
        | M.Free (e, pos) =>
          let val new = M.Free (cfExp e, pos)
          in  mkSomeList new
          end
        | M.Prop (e, pos) =>
          let val new = M.Prop (cfExp e, pos)
          in  mkSomeList new
          end
  
  and clean (l : M.t list option list) : M.t list =
      let fun f l =
              case l
               of [] => []
                | (NONE :: xs) => f xs
                | (SOME x :: xs) => x :: (f xs)
      in  List.concat (f l)
      end
      
  and cfStms (ss : M.t list) : M.t list = clean (List.map cfStm ss)
     
  fun cfTopStms stms = 
      let fun loop s =
              let val res = cfStms s
              in  if M.equalsStms (s, res)
                  then res
                  else (debugPrint "loop in stm optimization again"
                       ; loop res)
              end
      in  loop stms
      end
      
  fun cfFunction f =
      case f
       of F.T {name,
               rettype,
               args,
               locals,
               stm,
               pre,
               post,
               pos} => 
          F.T {name = name,
               rettype = rettype,
               args = args,
               locals = locals,
               stm = cfTopStms stm,
               pre = pre,
               post = post,
               pos = pos}
      
  fun cfProgram (p : P.t) : P.t =
      case p
       of P.T {structs,
               gvars,
               funcs,
               pos} =>
          P.T {structs = structs,
               gvars = gvars,
               funcs = List.map cfFunction funcs,
               pos = pos}
                       
  val top = cfProgram
end

  structure DeadCodeAfterReturn :
  sig
    val shrink : Ast.Program.t -> Ast.Program.t
  end =
  struct
    fun cfStm s =
        case s
         of M.AsnExp (lval, exp, pos) =>
          M.AsnExp (lval, exp, pos)
        | M.AsnCall (lval, id, es, pos) =>
          M.AsnCall (lval, id, es, pos)
        | M.AsnAlloc (lval, ty, pos) =>
          M.AsnAlloc (lval, ty, pos)
        | M.If (e, s1, s2, pos) =>
          M.If (e, cfStms s1, cfStms s2, pos)
        | M.While {invariant, condition, body, pos} =>
          M.While {invariant = invariant,
                   condition = condition,
                   body = cfStms body,
                   pos = pos}
        | M.Return (e, pos) =>
          M.Return (e, pos)
        | M.Call (id, es, pos) =>
          M.Call (id, es, pos)
        | M.Free (e, pos) =>
          M.Free (e, pos)
        | M.Prop (e, pos) =>
          M.Prop (e, pos)
    and cfStms ss =
        case ss
         of [] => []
          | x :: xs =>
            (case x
              of M.Return _ => [cfStm x]
               | _ => (cfStm x) :: (cfStms xs))
        
    fun shrinkStms ss =
        let fun loop s =
              let val res = cfStms s
              in  if M.equalsStms (s, res)
                  then res
                  else (debugPrint "loop in stm: dead code after return"
                       ; loop res)
              end
      in  loop ss
      end
    
    fun shrinkFun f =
        case f
       of F.T {name,
               rettype,
               args,
               locals,
               stm,
               pre,
               post,
               pos} => 
          F.T {name = name,
               rettype = rettype,
               args = args,
               locals = locals,
               stm = shrinkStms stm,
               pre = pre,
               post = post,
               pos = pos}
        
    fun shrink p =
        case p
       of P.T {structs,
               gvars,
               funcs,
               pos} =>
          P.T {structs = structs,
               gvars = gvars,
               funcs = List.map shrinkFun funcs,
               pos = pos}
  end 

  fun top ast = DeadCodeAfterReturn.shrink (Cf.top ast)
end
